import pymysql
import numpy as np
from matplotlib import pyplot as plt
from sklearn.linear_model import LinearRegression
from sklearn.tree import DecisionTreeRegressor


def xxhg():
    conn = pymysql.connect(host='127.0.0.1', port=3306, user='root', password='123456', db='citydata',
                           charset='utf8mb4', cursorclass=pymysql.cursors.DictCursor)
    cur = conn.cursor()

    sql = "SELECT * from population where type='年末常住人口数' and city='广东省' ORDER BY year"
    cur.execute(sql)
    data = cur.fetchall()
    print(data)

    # 提取特征和标签
    X = np.array([[d['year']] for d in data])
    y = np.array([d['number'] for d in data])

    # 创建并拟合模型
    model = LinearRegression().fit(X, y)

    # 进行预测
    pred_year = 2023  # 预测2023年的常住人口数
    predicted_number = model.predict([[pred_year]])
    print(f"预测{pred_year}年的常住人口数为：{predicted_number[0]:.2f}万人")


def jcs():
    conn = pymysql.connect(host='127.0.0.1', port=3306, user='root', password='123456', db='citydata',
                           charset='utf8mb4', cursorclass=pymysql.cursors.DictCursor)
    cur = conn.cursor()

    sql = "SELECT * from population where type='年末常住人口数' and city='广东省' ORDER BY year"
    cur.execute(sql)
    data = cur.fetchall()
    print(data)

    # 提取特征和标签
    X = np.array([[d['year']] for d in data])
    y = np.array([d['number'] for d in data])

    # 定义不同最大深度的决策树模型
    max_depths = [1, 2, 3, 4, 5]  # 尝试不同的最大深度
    models = {}
    for depth in max_depths:
        model = DecisionTreeRegressor(max_depth=depth)
        model.fit(X, y)
        print(model.predict([[2023]])[0])
        models[depth] = model

    # 绘制预测结果
    plt.figure()
    plt.scatter(X, y, s=20, edgecolor="black", c="darkorange", label="data")
    x_test = np.arange(2000, 2030, 1)[:, np.newaxis]
    for depth, model in models.items():
        y_1 = model.predict(x_test)
        plt.plot(x_test, y_1, label="max_depth=%d" % depth)

    plt.xlabel("year")
    plt.ylabel("population")
    plt.title("Decision Tree Regression")
    plt.legend()
    plt.show()


if __name__ == '__main__':
    xxhg()
    jcs()
